package com.qingcloud.df.executor.process.kube;

import com.google.common.collect.MoreCollectors;
import io.fabric8.kubernetes.api.model.ContainerStatus;
import io.fabric8.kubernetes.api.model.Pod;
import io.fabric8.kubernetes.client.informers.ResourceEventHandler;
import lombok.extern.slf4j.Slf4j;

import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Consumer;

@Slf4j
public class ExitCodeWatcher implements ResourceEventHandler<Pod> {

  private final String podName;
  private final String podNamespace;
  private final Consumer<Integer> onExitCode;
  private final Runnable onWatchFailure;

  private final AtomicBoolean active = new AtomicBoolean(true);

  public ExitCodeWatcher(final String podName,
                         final String podNamespace,
                         final Consumer<Integer> onExitCode,
                         final Runnable onWatchFailure) {
    this.podName = podName;
    this.podNamespace = podNamespace;
    this.onExitCode = onExitCode;
    this.onWatchFailure = onWatchFailure;
  }

  @Override
  public void onAdd(final Pod pod) {
    if (shouldCheckPod(pod)) {
      final Optional<Integer> exitCode = getExitCode(pod);
      exitCode.ifPresent(this::persistExitCode);
    }
  }

  @Override
  public void onUpdate(final Pod oldPod, final Pod newPod) {
    if (shouldCheckPod(newPod)) {
      final Optional<Integer> exitCode = getExitCode(newPod);
      exitCode.ifPresent(this::persistExitCode);
    }
  }

  @Override
  public void onDelete(final Pod pod, final boolean deletedFinalStateUnknown) {
    if (shouldCheckPod(pod)) {
      if (!deletedFinalStateUnknown) {
        final Optional<Integer> exitCode = getExitCode(pod);
        exitCode.ifPresentOrElse(
            this::persistExitCode,
            this::persistFailure);
      } else {
        persistFailure();
      }
    }
  }

  private boolean shouldCheckPod(final Pod pod) {
    final boolean correctName = podName.equals(pod.getMetadata().getName());
    final boolean correctNamespace = podNamespace.equals(pod.getMetadata().getNamespace());
    return active.get() && correctName && correctNamespace && KubePodResourceHelper.isTerminal(pod);
  }

  private Optional<Integer> getExitCode(final Pod pod) {
    final ContainerStatus mainContainerStatus = pod.getStatus().getContainerStatuses()
        .stream()
        .filter(containerStatus -> containerStatus.getName().equals(KubePodProcess.MAIN_CONTAINER_NAME))
        .collect(MoreCollectors.onlyElement());

    if (mainContainerStatus.getState() != null && mainContainerStatus.getState().getTerminated() != null) {
      return Optional.of(mainContainerStatus.getState().getTerminated().getExitCode());
    }
    return Optional.empty();
  }

  private void persistExitCode(final int exitCode) {
    if (active.compareAndSet(true, false)) {
      log.info("Received exit code {} for pod {}", exitCode, podName);
      onExitCode.accept(exitCode);
    }
  }

  private void persistFailure() {
    if (active.compareAndSet(true, false)) {
      log.error("Pod {} was deleted before we could retrieve its exit code", podName);
      onWatchFailure.run();
    }
  }

}
